package cn.njupt.tree;
//平衡二叉搜索树
public class AvlTree<E> extends BST<E> {

    public AvlTree(){
        super(null);
    }
    public AvlTree(Comparator<E> comparator){
        super(comparator);
    }

    /**
     * 判断以当前节点为根节点传递进来的子树是否平衡
     * @param node 当前节点
     * @return
     */
    public boolean  isBalance(Node<E> node){
        AvlNode<E> avlNode=((AvlNode<E>)node);
        return Math.abs(avlNode.balanceFactor())<=1;
    }

    public void updateHeight(Node<E> node){
        AvlNode<E> avlNode=((AvlNode<E>) node);
        avlNode.updateHeight();
    }


    static class AvlNode<E> extends Node{
        /**
         * AvlNode节点提供的属性，高度属性
         */
        int height=1;

        public AvlNode(Object element, BinaryTree.Node parent) {
            super(element, parent);
        }

        /**
         * 得到当前节点的平衡因子
         * @return
         */
        public int balanceFactor(){
            int left=this.left==null?0:((AvlNode<E>)(this.left)).height;
            int right=this.right==null?0:((AvlNode<E>)(this.right)).height;
            return left-right;
        }

        /**
         * 更新当前节点的高度
         * @return
         */
        public void updateHeight(){
            int leftHeight=this.left==null?0:((AvlNode<E>)(this.left)).height;
            int rightHeight=this.right==null?0:((AvlNode<E>)(this.right)).height;
            this.height=Math.max(leftHeight,rightHeight)+1;
        }

        public Node<E> tallerChildren(){
            int leftHeight=this.left==null?0:((AvlNode<E>)(this.left)).height;
            int rightHeight=this.right==null?0:((AvlNode<E>)(this.right)).height;
            if(leftHeight>rightHeight){
                return this.left;
            }else if(leftHeight<rightHeight){
                return this.right;
            }else{
                return this.isLeftChild()?this.left:this.right;
            }
        }


    }

    @Override
    public void afterAdd(Node<E> node) {
        while((node=node.parent)!=null){
            if(isBalance(node)){
                //如果当前节点平衡，就更新当前节点高度
                updateHeight(node);
            }else{
                //不平衡就进行恢复平衡
                reBalance2(node);
                break;
            }
        }
    }

    @Override
    public void afterRemove(Node<E> node) {
        while((node=node.parent)!=null){
            if(isBalance(node)){
                updateHeight(node);
            }else{
                reBalance(node);
            }
        }
    }

    /**
     * 恢复平衡，方法2
     * @param grand
     */
    private void reBalance2(Node<E> grand){
        Node<E> parent= ((AvlNode<E>) grand).tallerChildren();
        Node<E> child= ((AvlNode<E>) parent).tallerChildren();
        if(parent.isLeftChild()){
            if(child.isLeftChild()){
                //LL
                rotate(grand,child.left,child,child.right,parent,parent.right,grand,grand.right);
            }else{
                //LR
                rotate(grand,parent.left,parent,child.left,child,child.right,grand,grand.right);
            }
        }else{
            if(child.isRightChild()){
                //RR
                rotate(grand,grand.left,grand,parent.left,parent,child.left,child,child.right);
            }else{
                //RL
                rotate(grand,grand.left,grand,child.left,child,child.right,parent,parent.right);
            }
        }
    }

    /**
     * 统一进行旋转
     * @param grand  代表子树的根节点
     * @param a
     * @param b
     * @param c
     * @param d
     * @param e
     * @param f
     * @param g
     */
    public void rotate(Node<E> grand,
                       Node<E> a,Node<E> b,Node<E> c,
                       Node<E> d,
                       Node<E> e,Node<E> f,Node<E> g){
        b.left=a;
        b.right=c;

        f.left=e;
        f.right=g;

        d.left=b;
        d.right=f;
        d.parent=grand.parent;

        if(grand.parent==null){
            root=d;
        }else if(grand.isLeftChild()){
            grand.parent.left=d;
        }else{
            grand.parent.right=d;
        }

        if(a!=null){
            a.parent=b;
        }
        if(c!=null){
            c.parent=b;
        }
        if(e!=null){
            e.parent=f;
        }
        if(g!=null){
            g.parent=f;
        }
        b.parent=d;
        f.parent=d;
        updateHeight(f);
        updateHeight(b);
        updateHeight(d);
    }


    /**
     * 进行恢复平衡，方法1
     * @param grand 高度最低的那个不平衡节点
     */
    private void reBalance(Node<E> grand) {
        //传递过来的node节点显然是不平衡的节点
        //恢复平衡我们先要看是哪种类型，LL、LR、RR、RL
        Node<E> parent= ((AvlNode<E>) grand).tallerChildren();
        Node<E> child= ((AvlNode<E>) parent).tallerChildren();
        if(parent.isLeftChild()){
            if(child.isLeftChild()){
                //LL
                rotateRight(grand);
            }else{
                //LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        }else{
            if(child.isRightChild()){
                //RR
                rotateLeft(grand);
            }else{
                //RL
                rotateRight(parent);
                rotateLeft(grand);
            }
        }
    }

    /**
     * 左旋转
     * @param grand 传递的最低不平衡的那个节点
     */
    public void rotateLeft(Node<E> grand){
        Node<E> parent=grand.right;
        Node<E> child=parent.left;

        grand.right=child;

        parent.left=grand;

        if(grand.parent==null){
            root=parent;
        }else if(grand.isLeftChild()){
            grand.parent.left=parent;
        }else{
            grand.parent.right=parent;
        }
        parent.parent=grand.parent;
        grand.parent=parent;
        //child可能为null
        if(child!=null){
            child.parent=grand;
        }

        updateHeight(grand);
        updateHeight(parent);
    }

    /**
     * 右旋转
     * @param grand 传递的最低不平衡的那个节点
     */
    public void rotateRight(Node<E> grand){
        Node<E> parent=grand.left;
        Node<E> child=parent.right;

        grand.left=child;

        parent.right=grand;

        if(grand.parent==null){
            root=parent;
        }else if(grand.isLeftChild()){
            grand.parent.left=parent;
        }else{
            grand.parent.right=parent;
        }
        parent.parent=grand.parent;
        grand.parent=parent;
        //child可能为null
        if(child!=null){
            child.parent=grand;
        }

        updateHeight(grand);
        updateHeight(parent);
    }

    @Override
    public Node<E> createNode(E element, Node<E> parent) {
        return new AvlNode(element,parent);
    }
}
